Before submitting

• I agree to license this contribution under the terms of LICENSE.txt.
• I ran pre-commit run and committed any formatting changes. (clang-format n/a — Python only; black 26.3.1 -l 160 passes, file unchanged.)

Affected area

FE OSS kernels or CuTeDSL

Summary

compute-sanitizer --tool racecheck flags RAW hazards on kernel_gemm of
indexer_backward_sm100.py (IndexerBackwardSm100). Every flagged hazard traces to one
defect in _load_warp: all 32 lanes of the load warp store sW / sGradSignal to SMEM,
but only a single elected lane arrives on MBAR_W_LOADED:

# all 32 lanes store:
sGradSignal[pos] = mGradSignal[...]      # ceil(topk/32) stores per lane
sW[idx] = mW[...]                        # ceil(heads/32) stores per lane

cute.arch.fence_view_async_shared()      # proxy fence — does not order threads
with cute.arch.elect_one():
    cute.arch.mbarrier_arrive(mbar + MBAR_W_LOADED)   # one lane only

The compute warpgroup (warps 4–7) does mbarrier_wait(MBAR_W_LOADED) and then reads
sW / sGradSignal with plain shared loads. Per the PTX memory model, mbarrier.arrive
(default .release, .cta scope) releases only the prior memory accesses of the
executing thread, so the stores of the other 31 lanes have no happens-before edge to
the consumer's reads. This is a formal data race.

Fix (2 logical lines): initialize MBAR_W_LOADED with count WARP_SIZE and have all 32
lanes of the load warp arrive. Each lane's release-arrive then orders its own stores
before the consumer's acquire-wait. MBAR_W_LOADED is a one-shot handoff (single arrive
site, single phase-0 wait), so the count change is self-contained.

Why

Evidence that this is a real memory-model violation, not a sanitizer false positive:

1. Lane-0 exemption. A 1-CTA shape (b=1, s_q=1, s_kv=512, h=64, d=128, topk=128) with
   --racecheck-report hazard --print-limit 0 reports exactly 620 byte-level hazards:
   the recorded writer is always warp-0 lanes 1–31, and the same store instructions
   executed by lane 0 are never flagged. 620 = 4 sGradSignal stores × 31 lanes × 4 B +
   2 sW stores × 31 lanes × 2 B, byte-exact. Racecheck therefore models the mbarrier
   arrive→wait edge and applies it precisely to the arriving thread — matching the PTX
   model, and pinpointing exactly the 31 lanes whose stores are unordered.
2. No other site is flagged in the tested configurations. The other barriers in kernel_gemm — whole-group
   arrivals (DS_READY / DK_EMPTY / K_LOADED, count = 128) and single-signal
   tcgen05.commit barriers (S_FULL / DK_FULL / K_CONSUMED / DQ_DONE) — report zero
   hazards in the same runs, including at a multi-block shape (b=3, s_q=64, s_kv=2048, topk=512) that
   exercises the full 3-stage sK / 2-stage TMEM pipeline. W_LOADED was the only
   single-elected arrive standing in for whole-warp generic SMEM stores.
3. Standalone calibration (inlined below, plain CUDA, nvcc -arch=sm_100a):
   • V0, mirroring the kernel (32-lane store → lane-0-only arrive): 124 hazards
     = 31 lanes × 4 B, lane 0 exempt — the same signature as the real kernel.
   • V2 (all 32 lanes arrive, init count 32): 0 hazards.
   • V1 (__syncwarp() then lane-0-only arrive): still 124. This variant is arguably
     correct under the CUDA C++ __syncwarp memory-ordering guarantee, but racecheck does
     not chain the syncwarp edge into the mbarrier release; the all-lane arrive is the
     variant whose ordering is explicit in the PTX memory model and sanitizer-clean, so
     that is what this PR uses.

Why latent. We did not observe corruption on the build we tested (B200, CUDA 13.3,
nvidia-cutlass-dsl 4.6.1): d_index_q / d_weights are bitwise stable across 30 replays
even unpatched, and in the SASS we captured for the topk=128 kernel specialization, ptxas lowers the
release-arrive as an unpredicated MEMBAR.ALL.CTA executed by all lanes, which restores
the ordering as a side effect. That compensation is not contractual — an nvcc 13.3 compile of the same
pattern (calibration V0) contains no MEMBAR before the arrive — so the pattern should
not be relied on across compiler versions. Fixing it also makes kernel_gemm
racecheck-clean, so future real regressions are not buried in baseline noise.

Related issues

None filed; found by running compute-sanitizer over the DSA OSS kernels.

API and compatibility impact

No API change.

• Numerics: no difference observed on the tested shape/build. d_index_q / d_weights are byte-identical between unpatched
  and patched builds (direct byte comparison of outputs; full-sha256 stable 30/30 replays
  per side at b=3, s_q=64, s_kv=2048, topk=512). d_index_k is accumulated with fp32
  atomics and is run-to-run nondeterministic on both sides (maximum relative deviation from the first replay, over 30
  replays: 4.0e-6 unpatched vs 4.1e-6 patched).
• Performance: paired nsys kernel-time A/B on B200 (alternating unpatched/patched
  processes, 16 pairs per shape, 200 profiled launches per run, cuda_gpu_kern_sum):
  • b=3, s_q=64, s_kv=2048, topk=512: kernel_gemm −0.014% ± 0.017% (Student-t 95% CI
    on paired deltas) — no measurable change.
  • b=1, s_q=512, s_kv=512, topk=128: kernel_gemm +0.218% ± 0.016% (~+74 ns on a
    34.1 µs launch) — a small cost on this shape. It reproduces with the run
    order swapped (+0.204% ± 0.017%, 8 pairs), while the untouched ScoreGrad kernel
    measured in the same processes stays flat (|Δ| ≤ 0.03%), consistent with a small
    systematic cost rather than measurement noise. Mechanically the patch adds 31 extra
    single-thread arrives on one mbarrier, once per CTA, on the load warp.

Testing

Environment: B200 (SM100), CUDA 13.3, compute-sanitizer 2026.2.0.0, Nsight Systems
2026.3.1, nvidia-cutlass-dsl 4.6.1, torch 2.13, Python package built from this branch
(1.27.0).

racecheck (compute-sanitizer --tool racecheck, kernel driven through
indexer_backward_wrapper; "reports" = default analysis-mode error reports, which
aggregate the underlying per-byte hazards):

All unpatched reports point at the sW / sGradSignal stores in _load_warp; no other
site in the kernel is flagged before or after the patch.

pytest (test/python$ pytest fe_api/dsa/ -q, default L0): identical results on
unpatched base and this branch — 26 passed, 4 skipped, 4 failed, with the same 4 failing
cases on both sides (test_DSA_indexer_top_k ×2, test_DSA_sparse_attention_backward ×2
— reference-comparison failures present on unpatched develop@3a9ed3f in our environment,
untouched by this change). All test_DSA_indexer_backward cases pass on both sides.

Numerics and performance: as under "API and compatibility impact".

// Calibration: does racecheck model mbarrier arrive/wait edges, and for which threads?
// Pattern mirrors upstream indexer_backward_sm100 _load_warp:
//   producer warp0: all 32 lanes STS -> [optional sync] -> lane0-only mbarrier.arrive
//   consumer warp1: mbarrier wait -> LDS
// Variants (compile-time -DVARIANT=):
//   0: no intra-warp sync, lane0-only arrive   (upstream pattern)  -> 124 hazards, lane 0 exempt
//   1: __syncwarp() before lane0-only arrive                       -> 124 hazards
//   2: all-32-lane arrive (init count 32)                          -> 0 hazards
// Build/run: nvcc -arch=sm_100a -DVARIANT=0 mbar_calib.cu && compute-sanitizer --tool racecheck ./a.out
#include <cstdio>
#include <cuda_runtime.h>

#ifndef VARIANT
#define VARIANT 0
#endif

__device__ __forceinline__ void mbar_init(unsigned long long* mbar, unsigned count) {
  unsigned smaddr = (unsigned)__cvta_generic_to_shared(mbar);
  asm volatile("mbarrier.init.shared.b64 [%0], %1;" :: "r"(smaddr), "r"(count));
}
__device__ __forceinline__ void mbar_arrive(unsigned long long* mbar) {
  unsigned smaddr = (unsigned)__cvta_generic_to_shared(mbar);
  asm volatile("{.reg .b64 t; mbarrier.arrive.shared.b64 t, [%0];}" :: "r"(smaddr) : "memory");
}
__device__ __forceinline__ void mbar_wait(unsigned long long* mbar, unsigned phase) {
  unsigned smaddr = (unsigned)__cvta_generic_to_shared(mbar);
  asm volatile(
      "{.reg .pred P;\n"
      "WAIT_LOOP:\n"
      "mbarrier.try_wait.parity.shared.b64 P, [%0], %1;\n"
      "@!P bra WAIT_LOOP;}\n" :: "r"(smaddr), "r"(phase) : "memory");
}

__global__ void producer_consumer(float* out, const float* in) {
  __shared__ float buf[32];
  __shared__ unsigned long long mbar;
  int tid = threadIdx.x;
  int lane = tid % 32;
  int warp = tid / 32;
  if (tid == 0) {
#if VARIANT == 2
    mbar_init(&mbar, 32);
#else
    mbar_init(&mbar, 1);
#endif
  }
  __syncthreads();
  if (warp == 0) {
    buf[lane] = in[lane] * 2.0f;  // all 32 lanes store
#if VARIANT == 0
    if (lane == 0) mbar_arrive(&mbar);
#elif VARIANT == 1
    __syncwarp();
    if (lane == 0) mbar_arrive(&mbar);
#elif VARIANT == 2
    mbar_arrive(&mbar);
#endif
  } else if (warp == 1) {
    mbar_wait(&mbar, 0);
    out[lane] = buf[lane] + 1.0f;  // consumer reads
  }
}

int main() {
  float *d_in, *d_out;
  cudaMalloc(&d_in, 32 * sizeof(float));
  cudaMalloc(&d_out, 32 * sizeof(float));
  cudaMemset(d_in, 0, 32 * sizeof(float));
  producer_consumer<<<1, 64>>>(d_out, d_in);
  cudaError_t e = cudaDeviceSynchronize();
  printf("VARIANT=%d done err=%s\n", VARIANT, cudaGetErrorString(e));
  return 0;
}

Summary by CodeRabbit

• Bug Fixes
  • Improved synchronization for sparse attention gradient computation by adjusting how warp-level barriers are initialized and signaled.
  • Ensured proper ordering of shared-data writes before consumer waits to reduce potential timing/race issues and improve reliability.